import sys
import logging
import copy
import torch
from learners import get_learner
from utils.data_manager import DataManager
import os
from datetime import datetime
import shutil
import torch.backends
import torch.backends.cudnn
from tabulate import tabulate


def train(args):
    seed_list = copy.deepcopy(args["seed"])
    device = copy.deepcopy(args["device"])

    for seed in seed_list:
        args["seed"] = seed
        args["device"] = device
        if args.get("train_seed", None) is None:
            logging.info(
                "========== No train_seed in config, use seed instead. =========="
            )
            args["train_seed"] = args["seed"]
        _train(args)


def _train(args):

    # FIXME: This is a temporary solution to load the ViT models for fine-tuning
    assert isinstance(args["seed"], int), "seed must be an integer"
    # init_cls = 0 if args["init_cls"] == args["inc_cls"] else args["init_cls"]
    # assert init_cls == 0, "init_cls must be 0 for incremental learning"

    current_time = datetime.now().strftime("%Y%m%d%H%M%S")
    current_time = args["prefix"] if args["prefix"] else current_time
    logfile_path = "logs/{}/{}/{}".format(
        args["learner_name"],
        args["dataset"],
        current_time,
    )
    if not os.path.exists(logfile_path):
        os.makedirs(logfile_path)
    logfilename = os.path.join(
        logfile_path,
        "{}_{}_{}".format(
            args["prefix"],
            args["seed"],
            args["backbone_type"],
        ),
    )

    # 重置logging配置，确保不会被已有设置影响
    for handler in logging.root.handlers[:]:
        logging.root.removeHandler(handler)

    # 设置基本配置
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(filename)s] => %(message)s",
        handlers=[
            logging.FileHandler(filename=logfilename + ".log"),
            logging.StreamHandler(sys.stdout),
        ],
        force=True,  # 强制覆盖现有的logging配置
    )

    suffix = args["backbone_type"].split("_")[-1]
    shutil.copy2(args["config"], logfile_path)
    shutil.copy2(f"learners/{args['learner_name']}.py", logfile_path)
    try:
        shutil.copy2(f"backbone/vit_{suffix}.py", logfile_path)
    except FileNotFoundError:
        logging.info(
            f"backbone/vit_{suffix}.py not found, using backbone/vit.py instead."
        )
        shutil.copy2("backbone/vit.py", logfile_path)

    _set_random(args["train_seed"])
    _set_device(args)
    print_args(args)

    ckp_path = "_ckps/{}/{}/{}/{}".format(
        args["learner_name"],
        args["backbone_type"],
        args["dataset"],
        current_time,
    )
    args["ckp_path"] = ckp_path if getattr(args, "save_ckp", False) else None

    data_manager = DataManager(
        args["dataset"],
        args["shuffle"],
        args["seed"],
        args["init_cls"],
        args["inc_cls"],
        args,
    )

    args["nb_classes"] = data_manager.nb_classes  # update args
    args["nb_tasks"] = data_manager.nb_tasks
    model = get_learner(args["learner_name"], args, data_manager)

    cnn_curve, nme_curve = {"top1": [], "top5": []}, {
        "top1": [],
        "top5": [],
    }

    for task in range(data_manager.nb_tasks):
        # 记录任务标题
        print_task_header(task, data_manager.nb_tasks)

        # 模型训练
        cnn_accy, nme_accy = model.incremental_train()

        # 记录任务完成信息
        print_task_completion(task)

        if nme_accy is not None:
            # 更新准确率曲线
            cnn_curve["top1"].append(float(cnn_accy["top1"]))
            cnn_curve["top5"].append(float(cnn_accy["top5"]))
            nme_curve["top1"].append(float(nme_accy["top1"]))
            nme_curve["top5"].append(float(nme_accy["top5"]))

            print_accuracy_tables(cnn_accy, nme_accy)
            print_accuracy_curves(cnn_curve, nme_curve)
        else:
            # 更新准确率曲线
            cnn_curve["top1"].append(float(cnn_accy["top1"]))
            cnn_curve["top5"].append(float(cnn_accy["top5"]))

            print_accuracy_tables(cnn_accy)
            print_accuracy_curves(cnn_curve)

            # 记录没有NME准确率
            logging.info("No NME accuracy.")


def _set_device(args):
    device_type = args["device"]
    devices = []

    for device in device_type:
        if device == -1:
            device = torch.device("cpu")
        else:
            device = torch.device("cuda:{}".format(device))

        devices.append(device)

    args["device"] = devices


def _set_random(seed=1):
    logging.info("Setting training seed to {}".format(seed))
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


def print_args(args):
    for key, value in args.items():
        logging.info("{}: {}".format(key, value))


# 格式化输出辅助函数
def print_task_header(task_idx, total_tasks):
    """记录任务标题和分隔符"""
    separator = f"{'='*80}"
    task_header = f"Task {task_idx+1}/{total_tasks} - Training"

    logging.info(separator)
    logging.info(task_header)
    logging.info(separator)


def print_task_completion(task_idx):
    """记录任务完成信息"""
    separator = f"{'='*80}"
    logging.info(f"\nTraining Completed for Task {task_idx+1}")
    logging.info(separator)


def print_accuracy_tables(cnn_accy, nme_accy=None, log_raw=False):
    """记录准确率表格"""
    # 构建表头和数据
    has_nme = nme_accy is not None
    headers = ["Task/Class", "Top-1 (CNN)"]
    if has_nme:
        headers.append("Top-1 (NME)")
    table_data = []
    for key, value in cnn_accy["grouped"].items():
        if key != "total":
            row = [f"{key}", f"{float(value):.2f}%"]
            if has_nme:
                row.append(f"{float(nme_accy['grouped'][key]):.2f}%")
            table_data.append(row)
    logging.info(
        "\nAccuracy Table:"
        + "\n"
        + tabulate(table_data, headers=headers, tablefmt="pipe")
    )
    # 记录整体准确率 - 使用更清晰的格式
    overall_headers = ["Metric", "CNN", "NME"] if has_nme else ["Metric", "CNN"]
    overall_data = [
        (
            [
                "Top-1",
                f"{float(cnn_accy['top1']):.2f}%",
                f"{float(nme_accy['top1']):.2f}%",
            ]
            if has_nme
            else ["Top-1", f"{float(cnn_accy['top1']):.2f}%"]
        ),
        (
            [
                "Top-5",
                f"{float(cnn_accy['top5']):.2f}%",
                f"{float(nme_accy['top5']):.2f}%",
            ]
            if has_nme
            else ["Top-5", f"{float(cnn_accy['top5']):.2f}%"]
        ),
    ]
    logging.info(
        "\nOverall Accuracy:"
        + "\n"
        + tabulate(overall_data, headers=overall_headers, tablefmt="pipe")
    )

    # 记录原始数据
    if log_raw:
        logging.info(f"CNN: {cnn_accy['grouped']}")
        if has_nme:
            logging.info(f"NME: {nme_accy['grouped']}")


def print_accuracy_curves(cnn_curve, nme_curve=None, log_raw=False):
    """记录准确率曲线"""

    # 添加表头
    headers = ["Task", "CNN Top-1", "CNN Top-5"]
    if nme_curve is not None:
        headers.extend(["NME Top-1", "NME Top-5"])
    # 创建任务编号
    num_tasks = len(cnn_curve["top1"])
    # 格式化数据为表格形式 - 每个任务一行
    rows = []
    for i in range(num_tasks):
        row = [
            f"Task {i+1}",
            f"{float(cnn_curve['top1'][i]):.2f}%",
            f"{float(cnn_curve['top5'][i]):.2f}%",
        ]
        if nme_curve is not None:
            row.extend(
                [
                    f"{float(nme_curve['top1'][i]):.2f}%",
                    f"{float(nme_curve['top5'][i]):.2f}%",
                ]
            )
        rows.append(row)

    """计算平均精度"""
    top1_values = [float(v) for v in cnn_curve["top1"]]
    cnn_avg_top1 = sum(top1_values) / len(top1_values) if top1_values else 0
    top5_values = [float(v) for v in cnn_curve["top5"]]
    cnn_avg_top5 = sum(top5_values) / len(top5_values) if top5_values else 0

    if nme_curve is not None:
        # 确保所有值都是浮点数
        nme_top1_values = [float(v) for v in nme_curve["top1"]]
        nme_avg_top1 = (
            sum(nme_top1_values) / len(nme_top1_values) if nme_top1_values else 0
        )

        # 计算NME Top-5平均值
        nme_top5_values = [float(v) for v in nme_curve["top5"]]
        nme_avg_top5 = (
            sum(nme_top5_values) / len(nme_top5_values) if nme_top5_values else 0
        )
        avg_table = [
            [
                "Avg",
                f"{cnn_avg_top1:.2f}%",
                f"{cnn_avg_top5:.2f}%",
                f"{nme_avg_top1:.2f}%",
                f"{nme_avg_top5:.2f}%",
            ],
        ]
    else:
        nme_avg_top1 = None
        avg_table = [
            ["Avg", f"{cnn_avg_top1:.2f}%", f"{cnn_avg_top5:.2f}%"],
        ]

    rows.extend(avg_table)
    logging.info(
        "\nAccuracy Curves:" + "\n" + tabulate(rows, headers=headers, tablefmt="pipe")
    )

    # 记录原始格式 (保留完整数据的记录)
    if log_raw:
        logging.info(f"CNN top1 curve: {cnn_curve['top1']}")
        logging.info(f"CNN top5 curve: {cnn_curve['top5']}")
        if nme_curve is not None:
            logging.info(f"NME top1 curve: {nme_curve['top1']}")
            logging.info(f"NME top5 curve: {nme_curve['top5']}")


def print_average_accuracy(cnn_curve, nme_curve=None, log_raw=False):
    """记录平均准确率"""
    # 确保所有值都是浮点数
    top1_values = [float(v) for v in cnn_curve["top1"]]
    cnn_avg_top1 = sum(top1_values) / len(top1_values) if top1_values else 0

    # 计算Top-5平均值
    top5_values = [float(v) for v in cnn_curve["top5"]]
    cnn_avg_top5 = sum(top5_values) / len(top5_values) if top5_values else 0

    if nme_curve is not None:
        # 确保所有值都是浮点数
        nme_top1_values = [float(v) for v in nme_curve["top1"]]
        nme_avg_top1 = (
            sum(nme_top1_values) / len(nme_top1_values) if nme_top1_values else 0
        )

        # 计算NME Top-5平均值
        nme_top5_values = [float(v) for v in nme_curve["top5"]]
        nme_avg_top5 = (
            sum(nme_top5_values) / len(nme_top5_values) if nme_top5_values else 0
        )

        avg_table = [
            ["Metric", "CNN", "NME"],
            ["Avg Top-1", f"{cnn_avg_top1:.2f}%", f"{nme_avg_top1:.2f}%"],
            ["Avg Top-5", f"{cnn_avg_top5:.2f}%", f"{nme_avg_top5:.2f}%"],
        ]
    else:
        nme_avg_top1 = None
        nme_avg_top5 = None
        avg_table = [
            ["Metric", "CNN"],
            ["Avg Top-1", f"{cnn_avg_top1:.2f}%"],
            ["Avg Top-5", f"{cnn_avg_top5:.2f}%"],
        ]

    logging.info("\nAverage Accuracy:")
    logging.info("\n" + tabulate(avg_table, headers="firstrow", tablefmt="pipe"))

    # 同时记录原始格式
    if log_raw:
        logging.info(
            f"Average Accuracy (CNN): Top-1={cnn_avg_top1:.4f}, Top-5={cnn_avg_top5:.4f}"
        )
        if nme_curve is not None:
            logging.info(
                f"Average Accuracy (NME): Top-1={nme_avg_top1:.4f}, Top-5={nme_avg_top5:.4f}"
            )

    return cnn_avg_top1, nme_avg_top1
